import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
from sklearn.impute import SimpleImputer

# 初始化 Tushare Pro 接口
pro = ts.pro_api('276da5751ba9fda134012d591c37d0cbed93bdf1dad172617e265d6a')

# 定义股票代码列表
tickers = ['000001.SZ', '000002.SZ', '000003.SZ', '000004.SZ', '000005.SZ',
           '000006.SZ', '000007.SZ', '000008.SZ', '000009.SZ', '000010.SZ']

# 定义日期范围
start_date = '20230101'
end_date = '20250101'

# 获取股票数据
def get_stock_data(tickers, start_date, end_date):
    data = pd.DataFrame()
    for ticker in tickers:
        df = pro.daily(ts_code=ticker, start_date=start_date, end_date=end_date)
        df['ts_code'] = ticker  # 添加股票代码列
        data = pd.concat([data, df])
    return data

# 获取数据
stock_data = get_stock_data(tickers, start_date, end_date)
stock_data['trade_date'] = pd.to_datetime(stock_data['trade_date'])
stock_data.set_index('trade_date', inplace=True)
print(stock_data.head())


def calculate_technical_indicators(data):
    # 移动平均线
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA10'] = data['close'].rolling(window=10).mean()
    data['MA20'] = data['close'].rolling(window=20).mean()

    # 布林线
    data['BOLL_MA'] = data['close'].rolling(window=20).mean()
    data['BOLL_STD'] = data['close'].rolling(window=20).std()
    data['BOLL_UPPER'] = data['BOLL_MA'] + 2 * data['BOLL_STD']
    data['BOLL_LOWER'] = data['BOLL_MA'] - 2 * data['BOLL_STD']

    # RSI
    delta = data['close'].diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    RS = gain / loss
    data['RSI'] = 100 - (100 / (1 + RS))

    # MACD
    data['EMA12'] = data['close'].ewm(span=12, adjust=False).mean()
    data['EMA26'] = data['close'].ewm(span=26, adjust=False).mean()
    data['MACD'] = data['EMA12'] - data['EMA26']
    data['MACD_signal'] = data['MACD'].ewm(span=9, adjust=False).mean()

    # KDJ
    low_min = data['low'].rolling(window=9).min()
    high_max = data['high'].rolling(window=9).max()
    data['K'] = ((data['close'] - low_min) / (high_max - low_min) * 100).rolling(window=3).mean()
    data['D'] = data['K'].rolling(window=3).mean()
    data['J'] = 3 * data['K'] - 2 * data['D']

    # 其他指标
    data['ATR'] = data['high'].rolling(window=14).max() - data['low'].rolling(window=14).min()
    data['VWAP'] = (data['close'] * data['vol']).rolling(window=14).sum() / data['vol'].rolling(window=14).sum()
    data['BIAS'] = (data['close'] - data['MA10']) / data['MA10'] * 100
    data['OBV'] = data['vol'] * np.where(data['close'].diff() > 0, 1, -1).cumsum()
    data['W%R'] = (high_max - data['close']) / (high_max - low_min) * -100

    # 标签：根据收盘价涨跌
    data['Target'] = np.where(data['close'] > data['close'].shift(1), 1, 0)

    return data


# 计算技术指标
stock_data = calculate_technical_indicators(stock_data)
stock_data.dropna(inplace=True)  # 删除因计算指标产生的空值
print(stock_data.head())


def preprocess_data(data):
    # 删除无关列
    data.drop(columns=['ts_code', 'pre_close', 'change', 'pct_chg', 'amount'], inplace=True, errors='ignore')

    # 处理空值
    imputer = SimpleImputer(strategy='mean')
    data_imputed = pd.DataFrame(imputer.fit_transform(data), columns=data.columns, index=data.index)

    # 异常值处理（简单示例：删除超出3倍标准差的值）
    for col in data_imputed.columns:
        mean = data_imputed[col].mean()
        std = data_imputed[col].std()
        data_imputed = data_imputed[(data_imputed[col] > mean - 3 * std) & (data_imputed[col] < mean + 3 * std)]

    # 归一化
    scaler = MinMaxScaler()
    data_scaled = pd.DataFrame(scaler.fit_transform(data_imputed), columns=data_imputed.columns,
                               index=data_imputed.index)

    # 主成分分析
    pca = PCA(n_components=5)
    data_pca = pca.fit_transform(data_scaled)
    data_pca_df = pd.DataFrame(data_pca, index=data_scaled.index,
                               columns=[f'PC{i + 1}' for i in range(data_pca.shape[1])])

    # 相关性分析
    corr_matrix = data_scaled.corr()
    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    plt.show()

    return data_pca_df


# 数据预处理
preprocessed_data = preprocess_data(stock_data)
print(preprocessed_data.head())


def model_and_evaluate(data, target):
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)

    # 选择模型（逻辑回归）
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)

    # 模型评价
    y_pred = model.predict(X_test)
    print("Classification Report:")
    print(classification_report(y_test, y_pred))

    # 混淆矩阵
    cm = confusion_matrix(y_test, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()

    # ROC曲线
    y_prob = model.predict_proba(X_test)[:, 1]
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')

    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    import seaborn as sns
    from sklearn.preprocessing import StandardScaler, MinMaxScaler
    from sklearn.decomposition import PCA
    from sklearn.model_selection import train_test_split
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
    from sklearn.impute import SimpleImputer

    # 假设 stock_data 是已经计算了技术指标的 DataFrame
    # stock_data = ...

    # 1. 空值处理
    imputer = SimpleImputer(strategy='mean')
    stock_data_imputed = pd.DataFrame(imputer.fit_transform(stock_data), columns=stock_data.columns,
                                      index=stock_data.index)

    # 2. 异常值处理（简单示例：删除超出3倍标准差的值）
    for col in stock_data_imputed.columns:
        mean = stock_data_imputed[col].mean()
        std = stock_data_imputed[col].std()
        stock_data_imputed = stock_data_imputed[
            (stock_data_imputed[col] > mean - 3 * std) & (stock_data_imputed[col] < mean + 3 * std)]

    # 3. 归一化
    scaler = MinMaxScaler()
    stock_data_scaled = pd.DataFrame(scaler.fit_transform(stock_data_imputed), columns=stock_data_imputed.columns,
                                     index=stock_data_imputed.index)

    # 4. 主成分分析（PCA）
    pca = PCA(n_components=5)
    pca_result = pca.fit_transform(stock_data_scaled)
    pca_df = pd.DataFrame(pca_result, columns=[f'PC{i + 1}' for i in range(pca_result.shape[1])],
                          index=stock_data_scaled.index)

    # 5. 相关性分析
    corr_matrix = stock_data_scaled.corr()
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    plt.show()

    # 6. 个股画像（选择6个属性）
    # 假设选择以下6个属性：'close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD'
    selected_attributes = ['close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD']
    stock_data_selected = stock_data_scaled[selected_attributes]

    # 可视化个股画像
    for attribute in selected_attributes:
        plt.figure(figsize=(10, 6))
        sns.histplot(stock_data_selected[attribute], kde=True)
        plt.title(f'Distribution of {attribute}')
        plt.show()

    # 7. 数据均衡（跳过 SMOTE，直接使用原始数据）
    # 假设 'Target' 是目标列
    X = pca_df
    y = stock_data['Target']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 8. 使用逻辑回归模型进行建模与评价
    # 训练模型
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)

    # 模型预测
    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    # 模型评价
    print("Classification Report:")
    print(classification_report(y_test, y_pred))

    # 混淆矩阵
    cm = confusion_matrix(y_test, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()

    # ROC曲线
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    import seaborn as sns
    from sklearn.preprocessing import StandardScaler, MinMaxScaler
    from sklearn.decomposition import PCA
    from sklearn.model_selection import train_test_split
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
    from sklearn.impute import SimpleImputer

    # 假设 stock_data 是已经计算了技术指标的 DataFrame
    # stock_data = ...

    # 1. 空值处理
    imputer = SimpleImputer(strategy='mean')
    stock_data_imputed = pd.DataFrame(imputer.fit_transform(stock_data), columns=stock_data.columns,
                                      index=stock_data.index)

    # 2. 异常值处理（简单示例：删除超出3倍标准差的值）
    for col in stock_data_imputed.columns:
        mean = stock_data_imputed[col].mean()
        std = stock_data_imputed[col].std()
        stock_data_imputed = stock_data_imputed[
            (stock_data_imputed[col] > mean - 3 * std) & (stock_data_imputed[col] < mean + 3 * std)]

    # 3. 归一化
    scaler = MinMaxScaler()
    stock_data_scaled = pd.DataFrame(scaler.fit_transform(stock_data_imputed), columns=stock_data_imputed.columns,
                                     index=stock_data_imputed.index)

    # 4. 主成分分析（PCA）
    pca = PCA(n_components=5)
    pca_result = pca.fit_transform(stock_data_scaled)
    pca_df = pd.DataFrame(pca_result, columns=[f'PC{i + 1}' for i in range(pca_result.shape[1])],
                          index=stock_data_scaled.index)

    # 5. 相关性分析
    corr_matrix = stock_data_scaled.corr()
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    plt.show()

    # 6. 个股画像（选择6个属性）
    # 假设选择以下6个属性：'close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD'
    selected_attributes = ['close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD']
    stock_data_selected = stock_data_scaled[selected_attributes]

    # 可视化个股画像
    for attribute in selected_attributes:
        plt.figure(figsize=(10, 6))
        sns.histplot(stock_data_selected[attribute], kde=True)
        plt.title(f'Distribution of {attribute}')
        plt.show()

    # 7. 数据均衡（跳过 SMOTE，直接使用原始数据）
    # 假设 'Target' 是目标列
    X = pca_df
    y = stock_data['Target']

    # 确保 X 和 y 的索引一致
    X = X.loc[y.index]
    y = y.loc[X.index]

    # 检查 X 和 y 的形状
    print("Shape of X:", X.shape)
    print("Shape of y:", y.shape)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 8. 使用逻辑回归模型进行建模与评价
    # 训练模型
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)

    # 模型预测
    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    # 模型评价
    print("Classification Report:")
    print(classification_report(y_test, y_pred))

    # 混淆矩阵
    cm = confusion_matrix(y_test, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()

    # ROC曲线
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    import seaborn as sns
    from sklearn.preprocessing import StandardScaler, MinMaxScaler
    from sklearn.decomposition import PCA
    from sklearn.model_selection import train_test_split
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
    from sklearn.impute import SimpleImputer

    # 假设 stock_data 是已经计算了技术指标的 DataFrame
    # stock_data = ...

    # 1. 空值处理
    imputer = SimpleImputer(strategy='mean')
    stock_data_imputed = pd.DataFrame(imputer.fit_transform(stock_data), columns=stock_data.columns,
                                      index=stock_data.index)

    # 2. 异常值处理（简单示例：删除超出3倍标准差的值）
    for col in stock_data_imputed.columns:
        mean = stock_data_imputed[col].mean()
        std = stock_data_imputed[col].std()
        stock_data_imputed = stock_data_imputed[
            (stock_data_imputed[col] > mean - 3 * std) & (stock_data_imputed[col] < mean + 3 * std)]

    # 3. 归一化
    scaler = MinMaxScaler()
    stock_data_scaled = pd.DataFrame(scaler.fit_transform(stock_data_imputed), columns=stock_data_imputed.columns,
                                     index=stock_data_imputed.index)

    # 4. 主成分分析（PCA）
    pca = PCA(n_components=5)
    pca_result = pca.fit_transform(stock_data_scaled)
    pca_df = pd.DataFrame(pca_result, columns=[f'PC{i + 1}' for i in range(pca_result.shape[1])],
                          index=stock_data_scaled.index)

    # 5. 相关性分析
    corr_matrix = stock_data_scaled.corr()
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    plt.show()

    # 6. 个股画像（选择6个属性）
    # 假设选择以下6个属性：'close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD'
    selected_attributes = ['close', 'MA5', 'MA10', 'MA20', 'RSI', 'MACD']
    stock_data_selected = stock_data_scaled[selected_attributes]

    # 可视化个股画像
    for attribute in selected_attributes:
        plt.figure(figsize=(10, 6))
        sns.histplot(stock_data_selected[attribute], kde=True)
        plt.title(f'Distribution of {attribute}')
        plt.show()

    # 7. 数据均衡（跳过 SMOTE，直接使用原始数据）
    # 假设 'Target' 是目标列
    # 确保 'Target' 列在 stock_data 中
    if 'Target' not in stock_data.columns:
        stock_data['Target'] = np.where(stock_data_scaled['close'] > stock_data_scaled['close'].shift(1), 1, 0)
        stock_data.dropna(inplace=True)  # 删除因计算指标产生的空值

    X = stock_data_scaled.drop(columns=['Target'])
    y = stock_data['Target']

    # 检查 X 和 y 的形状
    print("Shape of X:", X.shape)
    print("Shape of y:", y.shape)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 8. 使用逻辑回归模型进行建模与评价
    # 训练模型
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)

    # 模型预测
    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    # 模型评价
    print("Classification Report:")
    print(classification_report(y_test, y_pred))

    # 混淆矩阵
    cm = confusion_matrix(y_test, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()

    # ROC曲线
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()
